from ipex_llm.transformers import AutoModelForCausalLM
from transformers import AutoTokenizer

model_path = "C:\models\deepseek-ai\DeepSeek-R1-Distill-Qwen-7B"
# model_path = "C:\models\deepseek-ai\DeepSeek-R1-Distill-Llama-8B"


tokenizer = AutoTokenizer.from_pretrained(model_path)

# When running LLMs on Intel iGPUs for Windows users, we recommend setting `cpu_embedding=True` in the from_pretrained function.
# This will allow the memory-intensive embedding layer to utilize the CPU instead of iGPU.
model_in_4bit = AutoModelForCausalLM.from_pretrained(pretrained_model_name_or_path=model_path,
                                                     load_in_4bit=True)
model_in_4bit_gpu = model_in_4bit.to('xpu')

# math_prompt = r"Please reason step by step, and put your final answer within \boxed{}."
math_prompt = r"请逐步推理，然后将最终答案放在\boxed {}中。"
prompt = math_prompt+"\n"+"今有雉兔同笼，上有三十五头，下有九十四足，问雉兔各几何？"

messages = [
    # {"role": "system", "content": "You are Qwen, created by Alibaba Cloud. You are a helpful assistant."},
    {"role": "user", "content": prompt}
]
text = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True
)
model_inputs = tokenizer([text], return_tensors="pt").to(model_in_4bit_gpu.device)

generated_ids = model_in_4bit_gpu.generate(
    **model_inputs,
    max_new_tokens=2048,
    temperature=0.6
)
generated_ids = [
    output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
]

response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
print(response)
import time
time.sleep(5)